# parser.py

import os
import json
from datetime import datetime
import requests
from sqlalchemy import create_engine, text

# ----------------------------
# 配置部分
# ----------------------------
DORIS_HOST = "192.168.37.1"
DORIS_PORT = 9030
DORIS_USER = "root"
DORIS_PASSWORD = "pass"
DATABASE = "demo"
TABLE = "production_line_equipment_usage"

JSON_DIR = "json"
INSERT_DIR = "insert"

os.makedirs(INSERT_DIR, exist_ok=True)

# 数据库配置
DB_CONFIG = {
    "user": DORIS_USER,
    "password": DORIS_PASSWORD,
    "host": DORIS_HOST,
    "port": DORIS_PORT
}

# 创建数据库引擎
def create_connection():
    try:
        connection_string = f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DATABASE}"
        engine = create_engine(connection_string, echo=False)

        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
            print("[INFO] 数据库连接成功")

        return engine
    except Exception as e:
        print(f"[ERROR] 数据库连接失败: {e}")
        return None


engine = create_connection()


def get_doris_table_columns():
    """
    使用 SQL 查询获取 Doris 表的字段列表
    """
    try:
        engine = create_engine(f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}"
                               f"@{DB_CONFIG['host']}:9030/{DATABASE}")

        with engine.connect() as conn:
            result = conn.execute(text("DESCRIBE fault_records"))
            columns = [row[0] for row in result if row[0] not in ['op_state', 'op_time']]
            print(f"[INFO] 成功获取表 `{TABLE}` 的字段列表：{columns}")
            return columns
    except Exception as e:
        print(f"[ERROR] 获取表结构失败：{e}")
        return []

def load_json_files():
    """
    加载 json/ 目录下的所有 JSON 文件
    """
    files = [f for f in os.listdir(JSON_DIR) if f.endswith('.json')]
    data_list = []
    for file in files:
        path = os.path.join(JSON_DIR, file)
        try:
            with open(path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                data_list.append((file, data))
        except Exception as e:
            print(f"[ERROR] 读取文件失败：{file} - {e}")
    print(f"[INFO] 已加载 {len(data_list)} 个 JSON 文件")
    return data_list


def process_data(data, valid_columns):
    """
    根据字段白名单处理数据，只保留有效字段，并添加默认字段（如 id、created_at）
    """
    # 提取结构化数据
    structured = data.get("structured", {})
    metadata = data.get("metadata", {})

    # 构造最终记录
    record = {}

    # 自动填充字段
    record['id'] = metadata.get('filename', '').replace('.json', '')
    record['created_at'] = metadata.get('record_time')

    # 只保留数据库中存在的字段
    for key in valid_columns:
        if key in ['id', 'created_at']:
            continue  # 跳过已手动设置的字段
        value = None
        if key in structured:
            value = structured[key]
        elif key.startswith('ik_'):
            value = data.get('ik_analysis', {}).get(key)
        record[key] = value

    return record


def insert_into_doris(engine, records):
    """
    将处理好的记录插入到 Doris 表中
    """
    if not engine or not records:
        print("[WARN] 没有可用数据或数据库连接失败，跳过插入")
        return

    try:
        with engine.connect() as conn:
            for record in records:
                columns = ', '.join(record.keys())
                placeholders = ', '.join([':' + k for k in record.keys()])
                stmt = text(f"INSERT INTO {TABLE} ({columns}) VALUES ({placeholders})")
                conn.execute(stmt, record)
            conn.commit()
        print(f"[INFO] 成功插入 {len(records)} 条记录到表 `{TABLE}`")
    except Exception as e:
        print(f"[ERROR] 插入失败：{e}")


def main():
    print("[INFO] 开始解析 JSON 数据并适配 Doris 表结构")

    # 获取 Doris 字段
    valid_columns = get_doris_table_columns()
    if not valid_columns:
        print("[ERROR] 没有可用字段，程序退出")
        return

    # 加载 JSON 数据
    json_data_list = load_json_files()

    # 处理每条数据
    count = 0
    all_records = []

    for filename, data in json_data_list:
        processed = process_data(data, valid_columns)

        output_filename = f"{os.path.splitext(filename)[0]}.json"
        output_path = os.path.join(INSERT_DIR, output_filename)

        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(processed, f, ensure_ascii=False, indent=4)
            count += 1
            print(f"[SAVE] 已写入：{output_path}")
            all_records.append(processed)
        except Exception as e:
            print(f"[ERROR] 写入文件失败：{output_path} - {e}")

    print(f"\n[INFO] 完成处理，共生成 {count} 条可用于插入的数据。\n")

    # 插入数据库
    insert_into_doris(engine, all_records)


if __name__ == "__main__":
    main()